{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load keypoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=6.37s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "from pycocotools.coco import COCO\n",
    "## install coco library https://github.com/pdollar/coco\n",
    "from collections import namedtuple\n",
    "import numpy as np\n",
    "from numpy import ma\n",
    "import math\n",
    "import skimage.io as io\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import pylab\n",
    "import json\n",
    "import time\n",
    "from numpy import linalg as LA\n",
    "pylab.rcParams['figure.figsize'] = (10.0, 8.0)\n",
    "\n",
    "## load keypoints json\n",
    "annFile = '/data/datasets/COCO/person_keypoints_trainval2014/person_keypoints_train2014.json' # keypoint file\n",
    "trainimagepath = '/data/guest_users/liangdong/liangdong/practice_demo/train2014/'             # train image path\n",
    "\n",
    "coco = COCO(annFile)\n",
    "cats = coco.loadCats(coco.getCatIds())\n",
    "catIds = coco.getCatIds(catNms=['person'])\n",
    "imgIds = coco.getImgIds(catIds=catIds )\n",
    "jointall=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "45174"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## train ids \n",
    "len(imgIds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transform data format from coco to real_time_pose "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "count = 0\n",
    "joint_all = dict()\n",
    "numimgs = len(imgIds)\n",
    "\n",
    "for i in range(0, numimgs):\n",
    "    #print('----image: '+str(imgIds[i])+' ---------------')\n",
    "    img = coco.loadImgs(imgIds[i])[0]\n",
    "    annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)\n",
    "    anns = coco.loadAnns(annIds)\n",
    "    \n",
    "    height = img['height']\n",
    "    width = img['width']\n",
    "    \n",
    "    prev_center = list(list())\n",
    "    segmentations = list()\n",
    "    \n",
    "    for j in range(len(anns)):\n",
    "        if anns[j]['num_keypoints'] < 5 or anns[j]['area'] < 32*32:\n",
    "            segmentations.append(anns[j]['segmentation'])\n",
    "            \n",
    "    for j in range(len(anns)):\n",
    "        \n",
    "        ## remove those persons whose keypoints are too small\n",
    "        if anns[j]['num_keypoints'] < 5 or anns[j]['area'] < 32*32:\n",
    "            continue\n",
    "        \n",
    "        person_center = [anns[j]['bbox'][0] + anns[j]['bbox'][2]/2,\n",
    "                         anns[j]['bbox'][1] + anns[j]['bbox'][3]/2]\n",
    "        flag = 0\n",
    "        isValidation = 0 \n",
    "        \n",
    "        for k in range(len(prev_center)):\n",
    "            dist1 = prev_center[k][0] - person_center[0]\n",
    "            dist2 = prev_center[k][1] - person_center[1]\n",
    "            #print dist1, dist2\n",
    "            if dist1*dist1+dist2*dist2 < prev_center[k][2]*0.3:\n",
    "                flag = 1\n",
    "                continue\n",
    "                \n",
    "        currentjoin={'isValidation': isValidation,\n",
    "                     'img_paths': trainimagepath + img['file_name'],\n",
    "                     'objpos': person_center,\n",
    "                     'image_id': img['id'],\n",
    "                     'bbox': anns[j]['bbox'],\n",
    "                     'img_width': width,\n",
    "                     'img_height': height,\n",
    "                     'segment_area': anns[j]['area'],\n",
    "                     'num_keypoints': anns[j]['num_keypoints'],\n",
    "                     'joint_self': np.zeros((17,3)).tolist(),\n",
    "                     'scale_provided': anns[j]['bbox'][3]/368.0,\n",
    "                     'segmentations': segmentations,\n",
    "                     'joint_others': {},\n",
    "                     'annolist_index': i ,\n",
    "                     'people_index': j,\n",
    "                     'numOtherPeople':0,\n",
    "                     'scale_provided_other':{},\n",
    "                     'objpos_other':{},\n",
    "                     'bbox_other':{},\n",
    "                     'segment_area_other':{},\n",
    "                     'num_keypoints_other':{}\n",
    "                    }    \n",
    "        \n",
    "        for part in range(17):\n",
    "            currentjoin['joint_self'][part][0] = anns[j]['keypoints'][part*3]\n",
    "            currentjoin['joint_self'][part][1] = anns[j]['keypoints'][part*3+1]\n",
    "            # 2 means cropped, 0 means occluded by still on image\n",
    "            if(anns[j]['keypoints'][part*3+2] == 2):\n",
    "                currentjoin['joint_self'][part][2] = 1\n",
    "            elif(anns[j]['keypoints'][part*3+2] == 1):\n",
    "                currentjoin['joint_self'][part][2] = 0                \n",
    "            else:\n",
    "                currentjoin['joint_self'][part][2] = 2\n",
    "            \n",
    "        count_other = 1     \n",
    "        currentjoin['joint_others'] ={}\n",
    "       \n",
    "        for k in range(len(anns)):\n",
    "            \n",
    "            if k==j or anns[k]['num_keypoints']==0:\n",
    "                continue\n",
    "                \n",
    "            annop = anns[k]\n",
    "            currentjoin['scale_provided_other'][count_other] = annop['bbox'][3]/368\n",
    "            currentjoin['objpos_other'][count_other] = [annop['bbox'][0]+annop['bbox'][2]/2, \n",
    "                                        annop['bbox'][1]+annop['bbox'][3]/2]\n",
    "            currentjoin['bbox_other'][count_other] = annop['bbox']\n",
    "            currentjoin['segment_area_other'][count_other] = annop['area']\n",
    "            currentjoin['num_keypoints_other'][count_other] = annop['num_keypoints']\n",
    "            currentjoin['joint_others'][count_other] = np.zeros((17,3)).tolist()\n",
    "            \n",
    "            for part in range(17):\n",
    "                currentjoin['joint_others'][count_other][part][0] = annop['keypoints'][part*3]\n",
    "                currentjoin['joint_others'][count_other][part][1] = annop['keypoints'][part*3+1]\n",
    "                \n",
    "                if(annop['keypoints'][part*3+2] == 2):\n",
    "                    currentjoin['joint_others'][count_other][part][2] = 1\n",
    "                elif(annop['keypoints'][part*3+2] == 1):\n",
    "                    currentjoin['joint_others'][count_other][part][2] = 0\n",
    "                else:\n",
    "                    currentjoin['joint_others'][count_other][part][2] = 2\n",
    "                \n",
    "              \n",
    "            currentjoin['numOtherPeople'] = len(currentjoin['joint_others']) \n",
    "            count_other = count_other + 1\n",
    "                        \n",
    "        prev_center.append([person_center[0], person_center[1],\n",
    "                            max(anns[j]['bbox'][2], anns[j]['bbox'][3])])\n",
    "        \n",
    "        joint_all[count] = currentjoin\n",
    "\n",
    "        count = count + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "valimagepath = '/data/guest_users/liangdong/liangdong/practice_demo/val2014/'\n",
    "valFlie = '/data/datasets/COCO/person_keypoints_trainval2014/person_keypoints_val2014.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=3.67s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "coco = COCO(valFlie)\n",
    "cats = coco.loadCats(coco.getCatIds())\n",
    "catIds = coco.getCatIds(catNms=['person'])\n",
    "imgIds = coco.getImgIds(catIds=catIds )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "21634"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(imgIds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "numimgs = len(imgIds)\n",
    "\n",
    "for i in range(0, numimgs):\n",
    "    #print('----image: '+str(imgIds[i])+' ---------------')\n",
    "    img = coco.loadImgs(imgIds[i])[0]\n",
    "    annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)\n",
    "    anns = coco.loadAnns(annIds)\n",
    "    \n",
    "    height = img['height']\n",
    "    width = img['width']\n",
    "    \n",
    "    prev_center = list(list())\n",
    "    segmentations = list()\n",
    "    \n",
    "    for j in range(len(anns)):\n",
    "        if anns[j]['num_keypoints'] < 5 or anns[j]['area'] < 32*32:\n",
    "            segmentations.append(anns[j]['segmentation'])\n",
    "            \n",
    "    for j in range(len(anns)):\n",
    "        \n",
    "        ## remove those persons whose keypoints are too small\n",
    "        if anns[j]['num_keypoints'] < 5 or anns[j]['area'] < 32*32:\n",
    "            continue\n",
    "        \n",
    "        person_center = [anns[j]['bbox'][0] + anns[j]['bbox'][2]/2,\n",
    "                         anns[j]['bbox'][1] + anns[j]['bbox'][3]/2]\n",
    "        flag = 0\n",
    "        isValidation = 0 \n",
    "        \n",
    "        for k in range(len(prev_center)):\n",
    "            dist1 = prev_center[k][0] - person_center[0]\n",
    "            dist2 = prev_center[k][1] - person_center[1]\n",
    "            #print dist1, dist2\n",
    "            if dist1*dist1+dist2*dist2 < prev_center[k][2]*0.3:\n",
    "                flag = 1\n",
    "                continue\n",
    "                \n",
    "        currentjoin={'isValidation': isValidation,\n",
    "                     'img_paths': valimagepath + img['file_name'],\n",
    "                     'objpos': person_center,\n",
    "                     'image_id': img['id'],\n",
    "                     'bbox': anns[j]['bbox'],\n",
    "                     'img_width': width,\n",
    "                     'img_height': height,\n",
    "                     'segment_area': anns[j]['area'],\n",
    "                     'num_keypoints': anns[j]['num_keypoints'],\n",
    "                     'joint_self': np.zeros((17,3)).tolist(),\n",
    "                     'scale_provided': anns[j]['bbox'][3]/368.0,\n",
    "                     'segmentations': segmentations,\n",
    "                     'joint_others': {},\n",
    "                     'annolist_index': i ,\n",
    "                     'people_index': j,\n",
    "                     'numOtherPeople':0,\n",
    "                     'scale_provided_other':{},\n",
    "                     'objpos_other':{},\n",
    "                     'bbox_other':{},\n",
    "                     'segment_area_other':{},\n",
    "                     'num_keypoints_other':{}\n",
    "                    }    \n",
    "        \n",
    "        for part in range(17):\n",
    "            currentjoin['joint_self'][part][0] = anns[j]['keypoints'][part*3]\n",
    "            currentjoin['joint_self'][part][1] = anns[j]['keypoints'][part*3+1]\n",
    "            # 2 means cropped, 0 means occluded by still on image\n",
    "            if(anns[j]['keypoints'][part*3+2] == 2):\n",
    "                currentjoin['joint_self'][part][2] = 1\n",
    "            elif(anns[j]['keypoints'][part*3+2] == 1):\n",
    "                currentjoin['joint_self'][part][2] = 0                \n",
    "            else:\n",
    "                currentjoin['joint_self'][part][2] = 2\n",
    "            \n",
    "        count_other = 1     \n",
    "        currentjoin['joint_others'] ={}\n",
    "       \n",
    "        for k in range(len(anns)):\n",
    "            \n",
    "            if k==j or anns[k]['num_keypoints']==0:\n",
    "                continue\n",
    "                \n",
    "            annop = anns[k]\n",
    "            currentjoin['scale_provided_other'][count_other] = annop['bbox'][3]/368\n",
    "            currentjoin['objpos_other'][count_other] = [annop['bbox'][0]+annop['bbox'][2]/2, \n",
    "                                        annop['bbox'][1]+annop['bbox'][3]/2]\n",
    "            currentjoin['bbox_other'][count_other] = annop['bbox']\n",
    "            currentjoin['segment_area_other'][count_other] = annop['area']\n",
    "            currentjoin['num_keypoints_other'][count_other] = annop['num_keypoints']\n",
    "            currentjoin['joint_others'][count_other] = np.zeros((17,3)).tolist()\n",
    "            \n",
    "            for part in range(17):\n",
    "                currentjoin['joint_others'][count_other][part][0] = annop['keypoints'][part*3]\n",
    "                currentjoin['joint_others'][count_other][part][1] = annop['keypoints'][part*3+1]\n",
    "                \n",
    "                if(annop['keypoints'][part*3+2] == 2):\n",
    "                    currentjoin['joint_others'][count_other][part][2] = 1\n",
    "                elif(annop['keypoints'][part*3+2] == 1):\n",
    "                    currentjoin['joint_others'][count_other][part][2] = 0\n",
    "                else:\n",
    "                    currentjoin['joint_others'][count_other][part][2] = 2\n",
    "                \n",
    "              \n",
    "            currentjoin['numOtherPeople'] = len(currentjoin['joint_others']) \n",
    "            count_other = count_other + 1\n",
    "                        \n",
    "        prev_center.append([person_center[0], person_center[1],\n",
    "                            max(anns[j]['bbox'][2], anns[j]['bbox'][3])])\n",
    "        \n",
    "        joint_all[count] = currentjoin\n",
    "\n",
    "        count = count + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## all joints length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "139861\n"
     ]
    }
   ],
   "source": [
    "print len(joint_all)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write json file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('trainval.json', 'w') as f:\n",
    "     json.dump(joint_all, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading data back"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('trainval.json', 'r') as f:\n",
    "     data = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
